#
#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
from datetime import datetime

from api.db import FileSource
from api.db.db_models import DB
from api.db.db_models import File, File2Document
from api.db.services.common_service import CommonService
from api.db.services.document_service import DocumentService
from api.utils import current_timestamp, datetime_format


# 定义文件与文档关联服务类，用于处理文件和文档之间的映射关系
class File2DocumentService(CommonService):
    # 指定该服务类操作的数据模型为File2Document
    model = File2Document

    @classmethod
    @DB.connection_context()
    def get_by_file_id(cls, file_id):
        # 根据文件ID获取所有关联的文档对象
        objs = cls.model.select().where(cls.model.file_id == file_id)
        return objs

    @classmethod
    @DB.connection_context()
    def get_by_document_id(cls, document_id):
        # 根据文档ID获取所有关联的文件对象
        objs = cls.model.select().where(cls.model.document_id == document_id)
        return objs

    @classmethod
    @DB.connection_context()
    def insert(cls, obj):
        # 插入一个新的文件-文档关联记录，若失败则抛出异常
        if not cls.save(**obj):
            raise RuntimeError("Database error (File)!")
        return File2Document(**obj)

    @classmethod
    @DB.connection_context()
    def delete_by_file_id(cls, file_id):
        # 根据文件ID删除对应的文件-文档关联记录
        return cls.model.delete().where(cls.model.file_id == file_id).execute()

    @classmethod
    @DB.connection_context()
    def delete_by_document_id(cls, doc_id):
        # 根据文档ID删除对应的文件-文档关联记录
        return cls.model.delete().where(cls.model.document_id == doc_id).execute()

    @classmethod
    @DB.connection_context()
    def update_by_file_id(cls, file_id, obj):
        # 根据文件ID更新文件-文档关联信息，并更新时间戳和日期
        obj["update_time"] = current_timestamp()
        obj["update_date"] = datetime_format(datetime.now())
        cls.model.update(obj).where(cls.model.id == file_id).execute()
        return File2Document(**obj)

    @classmethod
    @DB.connection_context()
    def get_storage_address(cls, doc_id=None, file_id=None):
        # 获取指定文档或文件的存储地址信息（kb_id 和 location）
        # 根据是否传入 doc_id 判断是通过文档还是文件查找存储路径
        if doc_id:
            f2d = cls.get_by_document_id(doc_id)
        else:
            f2d = cls.get_by_file_id(file_id)
        # 如果查到文件与文档的映射关系，则获取对应的文件对象
        if f2d:
            file = File.get_by_id(f2d[0].file_id)
            # 若文件来源为空或为本地类型，则返回文件父ID和存储路径
            if not file.source_type or file.source_type == FileSource.LOCAL:
                return file.parent_id, file.location
            # 否则更新 doc_id 为映射表中的文档 ID，用于后续查询文档位置
            doc_id = f2d[0].document_id

        # 确保最终有有效的 doc_id，否则抛出断言异常
        assert doc_id, "please specify doc_id"
        # 通过 DocumentService 查询文档详情
        e, doc = DocumentService.get_by_id(doc_id)
        # 返回知识库 ID 和文档的存储路径
        return doc.kb_id, doc.location
